import numpy as np
import matplotlib.pyplot as plt

size = 10#迭代次数
k = 4#聚类数
number = 200#粒子数

#需要聚类的数据集
train = np.random.random([number,2])

#4个类，K个类
center = np.zeros([k,2])

#4个类里面的具体点
center_train = [[] for i in range(k)]

colors = ["red","blue","orange","yellow","green"]

class Kmeans:
    def __init__(self,train,center,center_train,colors):
        self.train = train
        self.center = center
        self.center_train = center_train
        self.colors = colors

        #初始话聚类中心
        for i in range(center.shape[0]):
            tempX = np.random.choice(train.T[0])
            tempY = np.random.choice(train.T[1])
            if tempX not in center.T[0] and tempY not in center.T[1]:
                center[i] = np.array([tempX,tempY])

        #show
        plt.figure()
        plt.scatter(self.train.T[0],self.train.T[1],color="gray")
        plt.scatter(self.center.T[0],self.center.T[1],color=colors[0])
        plt.show()

    def update(self,size):
        #迭代size次
        for i in range(size):
            for j in range(self.train.shape[0]):
                distances = np.zeros(self.center.shape[0])
                for k in range(self.center.shape[0]):
                    distances[k] = (np.sqrt((self.train[j][0]-self.center[k][0])**2+(self.train[j][1]-self.center[k][1])**2))
                self.center_train[np.argmin(distances)].append(self.train[j])
            plt.title('iterations:' + str(i + 1))
            # print(self.center.shape[0])
            for i in range(self.center.shape[0]):
                plt.scatter(np.mat(self.center_train[i]).T[0].tolist()[0],np.mat(self.center_train[i]).T[1].tolist()[0] , color = self.colors[i+1])

            #更新聚类中心点
            self.center = np.zeros([self.center.shape[0],2])
            for i in range(self.center.shape[0]):
                tempX = np.mean(np.mat(self.center_train[i]).T[0])
                tempY = np.mean(np.mat(self.center_train[i]).T[1])
                self.center[i] = np.array([tempX,tempY])

            plt.scatter(center.T[0],center.T[1],color=colors[0])
            plt.show()

if __name__ == "__main__":
    km = Kmeans(train,center,center_train,colors)
    km.update(size)



